#!/bin/env python
#coding=utf-8
#authro : xl.hao<xiaoliang.hao@intcache.com>
#save ssh info, and auto login ssh dev

import pexpect, sqlite3
import sys, os
import base64
#import time

class c_ssh_dev:
	def __init__(self):
		self.name = ""
		self.ip   = ""
		self.port = 22
		self.user = "root"
		self.pswd = ""
		self.sups = ""

def sql_commit():
	con.commit()
def dev_query(key, cmd):
	print "%s : " % key
	ssh_dev = c_ssh_dev()
	wildcard_pos = key.find('%')
	if cmd == "get" or wildcard_pos != -1:
		sql = "select ip, port, user, password, devname, sups from intc_ssh_dev where devname like ? or ip like ?"
		csr.execute(sql, ('%'+key+'%', '%'+key+'%'))
	else:
		sql = "select ip, port, user, password, devname, sups from intc_ssh_dev where devname = ? COLLATE NOCASE or ip = ?"
		csr.execute(sql, (key, key))
	#print sql
	devs = csr.fetchall()
	count = 0
	for dev in devs:
		ssh_dev.ip, ssh_dev.port, ssh_dev.user, ssh_dev.pswd, ssh_dev.name, ssh_dev.sups = (dev[0], dev[1], dev[2], dev[3], dev[4], dev[5])
		if not ssh_dev.sups:
			ssh_dev.sups = ""
		print "dev %d : %s, ip : %s, port : %s, user : %s, sups : %d" % (count, ssh_dev.name, ssh_dev.ip, ssh_dev.port, ssh_dev.user, len(ssh_dev.sups))
		if count >= 3 and cmd == "no":
			break
		count = count + 1
	devs_len = len(devs)
	if devs_len > 1 and cmd == "no":
		print "device numbers(%d) greater than one" % (devs_len)
		lexit(1)
	elif devs_len == 0 and cmd == "no":
		print "no dev %s is found, please mail the dev info to xl.hao<xiaoliang.hao@intcache.com>" % key
		lexit(1)
	
	sys.stdout.flush()
	sys.stderr.flush()
	return ssh_dev

def get_dev(key, argv):
	if len(key) == 0:
		print "no dev_name or ip is assigned"
		usage()
		lexit(1)
	cmd = "no"
	sub_dev = ""
	if len(argv) >= 2:
		if argv[1] == "get":
			cmd = "get"
		elif argv[1] == "send":
			cmd = "no"
			if len(argv) < 3:
				print "send command need sub dev_name"
				usage()
				lexit(1)
			sub_dev = argv[2]
			if len(sub_dev) == 0:
				print "send command need sub dev_name"
				usage()
				lexit(1)

	ssh_dev = dev_query(key, cmd)
	ssh_dev2 = None
	if len(sub_dev) != 0:
		print "sub_dev : %s" % sub_dev
		ssh_dev2 = dev_query(sub_dev, cmd)
	return ssh_dev, ssh_dev2

def show_dev_all(devname):
	if len(devname) == 0:
		print "no dev_name or ip is assigned"
		usage()
		lexit(1)
	ssh_dev = dev_query(devname, "no")
	print "[%s %s %d %s %s %s]" % (ssh_dev.name, ssh_dev.ip, ssh_dev.port, ssh_dev.user, ssh_dev.pswd, ssh_dev.sups)


def create_db():
	con.execute('''CREATE TABLE intc_ssh_dev
		(ip CHAR(20)     NOT NULL,
		port          INT    DEFAULT(22),
		user          CHAR(60) NOT NULL,
		password      CHAR(512) NOT NULL,
		sups          CHAR(512),
		devname       CHAR(60)  PRIMARY KEY NOT NULL);''')
def base64_change():
	if os.path.exists("/var/dengludb/.intc_ssh_dev.b4"):
		rf = open("/var/dengludb/.intc_ssh_dev.b4", "rb")
		db4ct = base64.b64decode(rf.read())
		wf = open("/var/dengludb/.intc_ssh_dev.db", "wb")
		wf.write(db4ct)
		rf.close()
		wf.close()
	elif os.path.exists("/var/dengludb/.intc_ssh_dev.db"):
		rf = open("/var/dengludb/.intc_ssh_dev.db", "rb")
		eb4ct = base64.b64encode(rf.read())
		wf = open("/var/dengludb/.intc_ssh_dev.b4", "wb")
		wf.write(eb4ct)
		rf.close()
		wf.close()

def base64_dbchange():
	if os.path.exists("/var/dengludb/.intc_ssh_dev.db"):
		rf = open("/var/dengludb/.intc_ssh_dev.db", "rb")
		eb4ct = base64.b64encode(rf.read())
		wf = open("/var/dengludb/.intc_ssh_dev.b4", "wb")
		wf.write(eb4ct)
		rf.close()
		wf.close()
	elif os.path.exists("/var/dengludb/.intc_ssh_dev.b4"):
		rf = open("/var/dengludb/.intc_ssh_dev.b4", "rb")
		db4ct = base64.b64decode(rf.read())
		wf = open("/var/dengludb/.intc_ssh_dev.db", "wb")
		wf.write(db4ct)
		rf.close()
		wf.close()

def db_del():
	if os.path.exists("/var/dengludb/.intc_ssh_dev.db"):
		os.remove("/var/dengludb/.intc_ssh_dev.db")

def sql_connect():
	con = sqlite3.connect("/var/dengludb/.intc_ssh_dev.db")
	cusor = con.cursor()
	return con, cusor

def add_data(dev_info):
	devs = dev_info.split()
	if len(devs) < 5:
		print "%s info is lack" % devs
		lexit(1)
	if len(devs) > 5:
		print "%s info is more, add failed !!!" % devs
		lexit(1)
	print "%s" % devs
	try:
		sql = "insert into intc_ssh_dev (ip, port, user, password, devname) values (?, ?, ?, ?, ?)"
		csr.execute(sql, (devs[0], devs[1], devs[2], devs[3], devs[4]))
		sql_commit()
	except Exception as e:
		print e
		lexit(1)

def add_sups(dev, sups):
	print "%s" % dev
	try:
		sql = "update intc_ssh_dev set sups = ? where devname = ?"
		csr.execute(sql, (sups, dev))
		sql_commit()
	except Exception as e:
		print e
		lexit(1)

def alter():
	try:
		sql = "alter table intc_ssh_dev add sups CHAR(512)"
		csr.execute(sql)
		sql_commit()
	except Exception as e:
		print e
		lexit(1)

def update(dev_info):
	devs = dev_info.split()
	if len(devs) < 5:
		print "%s info is lack" % devs
		lexit(1)
	print "%s" % devs
	try:
		sql = "update intc_ssh_dev set ip = ?, port = ?, user = ?, password = ? where devname = ? COLLATE NOCASE"
		csr.execute(sql, (devs[0], devs[1], devs[2], devs[3], devs[4]))
		sql_commit()
	except Exception as e:
		print e
		lexit(1)

def del_data(dev_info):
	if len(dev_info) == 0:
		print "[%s] info is lack" % dev_info
		lexit(1)
	devs = dev_info.split()
	if len(devs[0]) == 0:
		print "%s info is lack" % devs
		lexit(1)
	print "del_data %s" % devs
	pswd = raw_input("Enter the delete passwd : ")
	if pswd != "haha_o_haha":
		print "delete password is not correct"
		lexit(1)
	try:
		sql = "delete from intc_ssh_dev  where devname = ?"
		csr.execute(sql, (devs[0], ))
		sql_commit()
	except Exception as e:
		print e
		lexit(1)

def usage():
	print "usage: "
	print "denglu <create_db|add|scp|fscp|addsu> [dev_name] <option>"

	print "sub command : "
	print "create_db                              : create db, now is disable"
	print "add 'ip port user password dev_name'   : add new dev info"
	print "up  'ip port user password dev_name'   : update dev info by dev_name"
	print "addsu dev_name 'password'              : add dev_name su (root) password, because some devices not allow root login"
	print "del dev_name                           : delete dev info by dev_name"
	print "scp  dev_name file1 file2              : scp file1, file2 ... to dev_name"
	print "fscp dev_name remote_file <local_file> : scp from dev_name's file to local dir(./),  if local_file is not given. Or else rename newname in local"
	print "batch dev_list_file command            : read device name from dev_list_file, and execute command on every device"
	print ""
	print "denglu device"
	print "dev_name                 : denglu dev"
	print "dev_name su              : denglu dev and su"
	print "key_name get             : get dev_names which include key_name. eg. denglu SDGD get"
	print "dev_name send dev_name2  : denglu dev_name2 through dev_name"
	print "dev_name -c cmd          : denglu dev, and execute cmd"
	lexit(0)

def check_ssh_key(key, kport):
	#login_user = os.getlogin()
	import getpass
	login_user = getpass.getuser()
	ssh_hosts_path = "/root/.ssh/known_hosts"
	if login_user == "root":
		ssh_hosts_path = "/root/.ssh/known_hosts"
	elif login_user == 'tmpuser':
		sys.exit(1)
	else:
		ssh_hosts_path = "/home/%s/.ssh/known_hosts" % login_user

	if kport != 22 :
		key = "[%s]:%d" % (key, kport)
	file = open(ssh_hosts_path)
	while 1:
		lines = file.readlines(10000)
		if not lines:
			break
		for line in lines:
			if line.find(key) != -1:
				file.close()
				return 1
	file.close()
	return 0

def getTerminalSize():
    import os
    env = os.environ
    def ioctl_GWINSZ(fd):
        try:
            import fcntl, termios, struct, os
            cr = struct.unpack('hh', fcntl.ioctl(fd, termios.TIOCGWINSZ, '1234'))
        except:
            return
        return cr
    cr = ioctl_GWINSZ(0) or ioctl_GWINSZ(1) or ioctl_GWINSZ(2)
    if not cr:
        try:
            print "os.ctermid"
            fd = os.open(os.ctermid(), os.O_RDONLY)
            cr = ioctl_GWINSZ(fd)
            os.close(fd)
        except:
            pass
    if not cr:
        print "env get"
        cr = (env.get('LINES', 25), env.get('COLUMNS', 80))

        ### Use get(key[, default]) instead of a try/catch
        #try:
        #    cr = (env['LINES'], env['COLUMNS'])
        #except:
        #    cr = (25, 80)
    return int(cr[0]), int(cr[1])

def scp(ssh_dev, files):
	str_files = ""
	if len(files) == 0:
		print "files is not given"
		lexit(1)

	#print files
	for file in files:
		if len(file) == 0:
			continue
		str_files = str_files + file + " "
	try:
		sspawn = "scp -rp -P %d %s %s@%s:/tmp/" % (ssh_dev.port, str_files, ssh_dev.user, ssh_dev.ip)
		print(sspawn)
		foo = pexpect.spawn(sspawn)
		#foo.expect('.*ssword:')
		ssh_newkey = 'Are you sure you want to continue connecting'
		pret = foo.expect([ssh_newkey, '.*ssword', '%'])
		if pret == 0:
			foo.sendline("yes")
			foo.expect('.*ssword')
			foo.sendline(ssh_dev.pswd)
		elif pret == 1:
			foo.sendline(ssh_dev.pswd)
		#elif pret == 2:
		#	foo.send("ok")
		#print(ssh_dev.pswd)
		foo.interact()
	except pexpect.EOF:
		pass
	except Exception as e:
		#print e
		pass

def fscp(ssh_dev, files):
	try:
		local_path = "./"
		file = ''
		files_len = len(files)
		if files_len == 0:
			print "file is not given"
			lexit(1)
		file = files[0]
		if files_len >= 2:
			if len(files[1]) != 0:
				local_path = "%s" % files[1]
		sspawn = "scp -rp -P %d %s@%s:%s %s" % (ssh_dev.port, ssh_dev.user, ssh_dev.ip, file, local_path)
		print (sspawn)
		print "send ssh"
		foo = pexpect.spawn(sspawn)
		print "wait ..."
		ssh_newkey = 'Are you sure you want to continue connecting'
		pret = foo.expect([ssh_newkey, '.*ssword', '%'])
		if pret == 0:
			foo.sendline("yes")
			foo.expect('.*ssword')
			foo.sendline(ssh_dev.pswd)
		elif pret == 1:
			print "send paswd"
			foo.sendline(ssh_dev.pswd)
		foo.interact()
	except pexpect.EOF:
		pass
	except Exception as e:
		#print e
		pass

def batchproc(devlist_file, batchcmd):
	import gevent
	devices = []
	if len(devlist_file) == 0:
		print "devlist_file need"
		usage()
		lexit(1)
	try:
		fd = open(devlist_file, "rb")
		lines = fd.readlines()
		for line in lines:
			if len(line) == 0:
				continue
			if line.startswith('#'):
				continue
			line = line.strip("\r\n")
			print "read %s" % line
			devices.append(line)
		fd.close()
	except Exception as e:
		print e
		lexit(1)

	tasks = []
	newcmd = batchcmd
	#batchoutfile = "batch.log"

	if len(newcmd) == 0:
		print "command is null"
		lexit(1)

	#batchoutfd = open(batchoutfile, "wb")
	for device in devices:
		ssh_dev, tmp = get_dev(device, "")
		_g = gevent.spawn(commandexec, ssh_dev, newcmd)
		tasks.append(_g)
	db_del()
	for task in tasks:
		task.join()

def commandexec(ssh_dev, cmd):
	try:
		ssh_newkey = 'Are you sure you want to continue connecting'
		sspawn = "ssh -p %d -l %s %s %s" % (ssh_dev.port, ssh_dev.user, ssh_dev.ip, cmd)
		foo = pexpect.spawn(sspawn)
		#foo.logfile_read = sys.stdout
		#print(ssh_dev.pswd)
		#pret = foo.expect([ssh_newkey, '.*ssword', '].*'])
		pret = foo.expect([ssh_newkey, '.*ssword', pexpect.EOF])
		if pret == 0:
			foo.sendline("yes")
			foo.expect('.*ssword')
			foo.sendline(ssh_dev.pswd)
			foo.expect(pexpect.EOF)
		elif pret == 1:
			foo.sendline(ssh_dev.pswd)
			foo.expect(pexpect.EOF)
		#elif pret == 2:
		#foo.sendline(ssh_dev.pswd)
		#foo.expect(pexpect.EOF)
		print('\033[1;31m')
		print "%s : " % ssh_dev.name
		print foo.before
		print('\033[0m')
	except Exception as e:
		print('\033[1;33m')
		print "%s : " % ssh_dev.name
		print e
		print('\033[0m')
		pass

def denglu(ssh_dev, ssh_dev2, cmd):
	try:
		(row, col) = getTerminalSize()
		#input_yes = False
		ssh_newkey = 'Are you sure you want to continue connecting'
		sspawn = "ssh -o ConnectTimeout=30 -p %d %s@%s" % (ssh_dev.port, ssh_dev.user, ssh_dev.ip)
		#if check_ssh_key(ssh_dev.ip, ssh_dev.port) == 0:
		#	input_yes = True
		foo = pexpect.spawn(sspawn)
		#if input_yes:
		#	foo.expect('(.*yes/no*|.*ssword)')
		#	foo.sendline("yes")
		#print(ssh_dev.pswd)
		#pret = foo.expect([ssh_newkey, '.*ssword', ']($|#)'])
		pret = foo.expect([ssh_newkey, '.*ssword:', ']($|#)'])
		#print "pret : %s, %s" % (pret, foo.after)
		if pret == 0:
			foo.sendline("yes")
			foo.expect('.*ssword:')
			foo.sendline(ssh_dev.pswd)
		elif pret == 1:
			foo.sendline(ssh_dev.pswd)
		elif pret == 2:
			foo.sendline('alias grep="grep --color=auto" && export TMOUT=900')
		if ssh_dev2:
			#foo.expect(']($|#)')
			foo.expect(']')
			sspawn2 = "ssh -o ConnectTimeout=30 -p %d %s@%s" % (ssh_dev2.port, ssh_dev2.user, ssh_dev2.ip)
			print(sspawn2)
			foo.sendline(sspawn2)
			#sub_pret = foo.expect(['.*ssword:', ']($|#)'])
			sub_pret = foo.expect(['.*ssword:', ']($|#)', ssh_newkey])
			print "sub_pret %d" % sub_pret
			if sub_pret == 0:
				foo.sendline(ssh_dev2.pswd)
			elif sub_pret == 1:
				foo.sendline('alias grep="grep --color=auto" && export TMOUT=900')
			elif sub_pret == 2:
				foo.sendline("yes")
				sub_pret = foo.expect('.*ssword:')
				foo.sendline(ssh_dev2.pswd)
		pret = foo.expect(['.*ssword:', ']'])
		#print (foo.before, foo.after)
		if pret == 0:
			print "%s psword maybe error, please update %s psword!" % (ssh_dev.name, ssh_dev.name)
			foo.sendline("\r")
			foo.interact()
			return
		elif pret == 1:
			if len(ssh_dev.sups) != 0 and ssh_dev.user != "root":
				sspawn2 = "su -"
				print(sspawn2)
				foo.sendline(sspawn2)
				foo.expect('.*ssword:')
				foo.sendline(ssh_dev.sups)
				#_pret = foo.expect(['.*ssword', ']($|#)'])
				#if _pret == 0:
				#	print "%s su psword maybe error, please update %s su psword!" % (ssh_dev.name, ssh_dev.name)
				#elif _pret == 1:
				#	foo.sendline('alias grep="grep --color=auto" && export TMOUT=900')
			foo.sendline('alias grep="grep --color=auto" && export TMOUT=900')
		if len(cmd) != 0:
			#print "send %s" % cmd
			foo.expect(']($|#)')
			foo.sendline(cmd)
		foo.setwinsize(row, col)
		#rows, cols = foo.getwinsize()
		#print (rows, cols)
		foo.interact()
	except pexpect.EOF as e1:
		if foo.before is not None:
			print foo.before
		else:
			print e1
		pass
	except Exception as e:
		print e
		pass
def lexit(d):
	db_del()
	os._exit(d)

if __name__ == '__main__':
	base64_change()
	con, csr = sql_connect()
	cmd = ""

	argv_len = len(sys.argv)
	if argv_len > 1 :
		if sys.argv[1] == "create_db":
			#create_db()
			#base64_dbchange()
			#db_del()
			lexit(0)
		elif sys.argv[1] == "alter":
			#alter()
			#base64_dbchange()
			#db_del()
			lexit(0)
		elif sys.argv[1] == "add":
			if argv_len < 3:
				print "dev info is need"
				usage()
			else:
				add_data(sys.argv[2])
				base64_dbchange()
				db_del()
				lexit(0)
		elif sys.argv[1] == "up":
			if argv_len < 3:
				print "dev info is need"
				usage()
			else:
				update(sys.argv[2])
				base64_dbchange()
				db_del()
				lexit(0)
		elif sys.argv[1] == "addsu":
			if argv_len < 4:
				print "add su password info error"
				usage()
			else:
				add_sups(sys.argv[2], sys.argv[3])
				base64_dbchange()
				db_del()
				lexit(0)
		elif sys.argv[1] == "del":
			if argv_len < 3:
				print "dev_name is need"
				usage()
			else:
				del_data(sys.argv[2])
				base64_dbchange()
				db_del()
				lexit(0)
		elif sys.argv[1] == "help" or sys.argv[1] == "--help" or sys.argv[1] == "-h":
			usage()
		elif sys.argv[1] == "scp":
			if argv_len < 4:
				print "files is need"
				usage()
			else:
				ssh_dev, tmp = get_dev(sys.argv[2], sys.argv[3:])
				db_del()
				scp(ssh_dev, sys.argv[3:])
				lexit(0)
		elif sys.argv[1] == "fscp":
			if argv_len < 4:
				print "file is need"
				usage()
			else:
				ssh_dev, tmp = get_dev(sys.argv[2], sys.argv[3:])
				db_del()
				fscp(ssh_dev, sys.argv[3:])
				lexit(0)
		elif sys.argv[1] == "ppswdd":
			if argv_len < 3:
				usage()
			else:
				show_dev_all(sys.argv[2])
				db_del()
				lexit(0)
		elif sys.argv[1] == "batch":
			if argv_len < 4:
				usage()
			else:
				batchproc(sys.argv[2], sys.argv[3])
				lexit(0)
		else:
			ssh_dev, ssh_dev2 = get_dev(sys.argv[1], sys.argv[1:])
			db_del()
			if argv_len > 2:
				if sys.argv[2] == '-c':
					if argv_len < 4:
						print "sub command is need"
						usage()
					else:
						cmd = sys.argv[3]
				if sys.argv[2] == 'su':
						cmd = ""
				elif sys.argv[2] == "get":
					print "get %s devives =======================>" % (sys.argv[1])
					sys.stdout.flush()
					lexit(0)
			denglu(ssh_dev, ssh_dev2, cmd)
